#include "PBurgProcessor.h"
#include <QDebug>

// Main entry function
PBurgProcessor::SpectrumResult PBurgProcessor::pburg(
    const QVector<double>& x, int order, double samplingFreq, int nfft, const QString& range)
{
    // 参数校验
    if (order <= 0 || order >= x.size())  {
        qWarning() << "Invalid order value";
        return {};
    }

    // Burg算法计算AR参数
    QVector<double> a, e, k;
    if (!BurgAlgorithm::compute(order, x, a, e, k)) {
        return {};
    }

    // 计算功率谱
    QVector<double> freqBins, psd;
    SpectrumCalculator::computePSD(a, e.last(),  nfft, samplingFreq, freqBins, psd);

    // 频率范围处理
    if (range.toLower()  == "whole") {
        // 完整频谱处理逻辑
        const int half = nfft / 2 + 1;
        freqBins = Utils::linspace(0, samplingFreq, nfft).mid(0, nfft);
        QVector<double> psdWhole(nfft);
        std::copy(psd.begin(),  psd.end(),  psdWhole.begin());
        std::reverse_copy(psd.begin()  + 1, psd.end()  - 1, psdWhole.begin()  + half);
        return {freqBins, psdWhole};
    }
    return {freqBins, psd};
}

// Burg算法实现
bool PBurgProcessor::BurgAlgorithm::compute(
    int order, const QVector<double>& x, QVector<double>& a,
    QVector<double>& e, QVector<double>& k)
{
    const int N = x.size();
    QVector<double> ef = x.toList().mid(1);   // 前向预测误差
    QVector<double> eb = x.toList().mid(0,  N-1);  // 后向预测误差

    e.resize(order  + 1);
    k.resize(order);
    e[0] = x.sliced(0).toVector().squaredSum()  / N;

    for (int m = 0; m < order; ++m) {
        // 计算反射系数
        double numerator = 0.0, denominator = 0.0;
        for (int i = 0; i < N - m - 1; ++i) {
            numerator += ef[i] * eb[i];
            denominator += ef[i] * ef[i] + eb[i] * eb[i];
        }
        k[m] = 2 * numerator / denominator;

        // 更新预测误差
        for (int i = 0; i < N - m - 1; ++i) {
            double tmp = ef[i];
            ef[i] -= k[m] * eb[i];
            eb[i] = eb[i] - k[m] * tmp;
        }

        // 更新AR系数
        if (m == 0) {
            a.resize(2);
            a[0] = 1.0;
            a[1] = -k[m];
        } else {
            QVector<double> aNew(a.size()  + 1);
            aNew[0] = 1.0;
            for (int j = 1; j <= m; ++j) {
                aNew[j] = a[j] - k[m] * a[m - j];
            }
            aNew[m + 1] = -k[m];
            a = aNew;
        }

        e[m + 1] = e[m] * (1 - k[m] * k[m]);
    }
    return true;
}

// 功率谱计算
void PBurgProcessor::SpectrumCalculator::computePSD(
    const QVector<double>& a, double e, int nfft, double samplingFreq,
    QVector<double>& freqBins, QVector<double>& psd)
{
    // 计算频率响应
    QVector<std::complex<double>> h(nfft, {0,0});
    for (int k = 0; k < a.size();  ++k) {
        h[k] = a[k];
    }
    QVector<std::complex<double>> H = Utils::fft(h, nfft);

    // 计算PSD
    const double Ts = 1.0 / samplingFreq;
    freqBins = Utils::linspace(0, samplingFreq/2, nfft/2 + 1);
    psd.resize(freqBins.size());
    for (int i = 0; i < psd.size();  ++i) {
        psd[i] = e * Ts / std::norm(H[i]);
    }
}

// 辅助函数实现
QVector<double> PBurgProcessor::Utils::linspace(double start, double end, int num) {
    QVector<double> vec(num);
    double step = (end - start) / (num - 1);
    for (int i = 0; i < num; ++i) {
        vec[i] = start + i * step;
    }
    return vec;
}

QVector<std::complex<double>> PBurgProcessor::Utils::fft(const QVector<double>& x, int nfft) {
    // 此处可集成用户之前提供的FFT实现
    QVector<std::complex<double>> result(nfft);
    // ... (使用用户提供的FFT代码)
    return result;
}
